"""
一键优化补丁 - 直接替换 run_interactive.py 中的 stage_5 和 stage_6

使用方法:
1. 将此文件保存为 mm_story_agent/tools/optimized_image_generator.py
2. 在 run_interactive.py 中导入并使用新方法
"""

import asyncio
import aiohttp
import time
import hashlib
import json
from pathlib import Path
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass
import sys


@dataclass
class ImageGenerationResult:
    """图像生成结果"""
    index: int
    path: Optional[str]
    success: bool
    time: float
    cached: bool = False
    error: Optional[str] = None


class OptimizedImageGenerator:
    """优化的图像生成器 - 支持并行生成和智能缓存"""
    
    def __init__(self, 
                 api_key: str,
                 max_concurrent: int = 4,
                 request_interval: float = 1.0,
                 cache_dir: Optional[Path] = None,
                 max_retries: int = 3):
        """
        Args:
            api_key: DashScope API密钥
            max_concurrent: 最大并发数（推荐2-4）
            request_interval: 请求间隔秒数
            cache_dir: 缓存目录
            max_retries: 失败重试次数
        """
        self.api_key = api_key
        self.max_concurrent = max_concurrent
        self.request_interval = request_interval
        self.cache_dir = cache_dir
        self.max_retries = max_retries
        
        # 统计信息
        self.stats = {
            'total': 0,
            'success': 0,
            'failed': 0,
            'cached': 0,
            'retry_count': 0,
            'total_time': 0
        }
        
        # 进度显示
        self.completed = 0
        
        if cache_dir:
            cache_dir.mkdir(parents=True, exist_ok=True)
    
    def _get_cache_key(self, prompt: str) -> str:
        """生成缓存键"""
        return hashlib.md5(prompt.encode('utf-8')).hexdigest()
    
    def _check_cache(self, cache_key: str) -> Optional[str]:
        """检查缓存"""
        if not self.cache_dir:
            return None
        
        cache_file = self.cache_dir / f"{cache_key}.json"
        if cache_file.exists():
            try:
                with open(cache_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    image_path = Path(data.get('image_path', ''))
                    if image_path.exists():
                        return str(image_path)
            except:
                pass
        return None
    
    def _save_cache(self, cache_key: str, image_path: str, prompt: str):
        """保存缓存"""
        if not self.cache_dir:
            return
        
        cache_file = self.cache_dir / f"{cache_key}.json"
        try:
            with open(cache_file, 'w', encoding='utf-8') as f:
                json.dump({
                    'image_path': image_path,
                    'prompt': prompt,
                    'timestamp': time.time()
                }, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"⚠️  缓存保存失败: {e}")
    
    async def _wait_for_task(self, 
                             session: aiohttp.ClientSession,
                             task_id: str,
                             max_wait: int = 180) -> Optional[str]:
        """轮询任务直到完成（使用自适应轮询策略）"""
        from mm_story_agent.performance_optimizer import AdaptiveTaskPoller
        
        poller = AdaptiveTaskPoller(
            initial_interval=1.0,  # 初始1秒
            min_interval=0.5,      # 最快0.5秒
            max_interval=5.0       # 最慢5秒
        )
        
        return await poller.poll_task(session, task_id, self.api_key, max_wait)
    
    async def _download_image(self,
                             session: aiohttp.ClientSession,
                             url: str,
                             save_path: Path,
                             index: int) -> str:
        """下载图片"""
        image_path = save_path / f"p{index:04d}.png"
        
        try:
            async with session.get(url, timeout=30) as resp:
                if resp.status == 200:
                    content = await resp.read()
                    with open(image_path, 'wb') as f:
                        f.write(content)
                    return str(image_path)
                else:
                    raise Exception(f"下载失败: HTTP {resp.status}")
        except Exception as e:
            raise Exception(f"下载图片失败: {e}")
    
    async def _generate_single(self,
                               session: aiohttp.ClientSession,
                               semaphore: asyncio.Semaphore,
                               prompt: str,
                               save_path: Path,
                               index: int) -> ImageGenerationResult:
        """生成单张图片"""
        async with semaphore:  # 控制并发
            start_time = time.time()
            
            # 检查缓存
            cache_key = self._get_cache_key(prompt)
            cached_path = self._check_cache(cache_key)
            if cached_path:
                self.stats['cached'] += 1
                self._update_progress(index, cached=True)
                return ImageGenerationResult(
                    index=index,
                    path=cached_path,
                    success=True,
                    time=0,
                    cached=True
                )
            
            # API请求
            url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "X-DashScope-Async": "enable"
            }
            payload = {
                "model": "wanx-v1",
                "input": {"prompt": prompt},
                "parameters": {"size": "1024*1024", "n": 1}
            }
            
            # 重试机制
            last_error = None
            for attempt in range(self.max_retries):
                try:
                    # 提交任务
                    async with session.post(url, json=payload, headers=headers, timeout=30) as resp:
                        if resp.status == 200:
                            result = await resp.json()
                            task_id = result.get('output', {}).get('task_id')
                            
                            if task_id:
                                # 等待任务完成
                                image_url = await self._wait_for_task(session, task_id)
                                
                                # 下载图片
                                image_path = await self._download_image(
                                    session, image_url, save_path, index
                                )
                                
                                # 保存缓存
                                self._save_cache(cache_key, image_path, prompt)
                                
                                elapsed = time.time() - start_time
                                self.stats['success'] += 1
                                self.stats['total_time'] += elapsed
                                
                                self._update_progress(index, success=True, time=elapsed)
                                
                                return ImageGenerationResult(
                                    index=index,
                                    path=image_path,
                                    success=True,
                                    time=elapsed
                                )
                        else:
                            error_text = await resp.text()
                            raise Exception(f"API错误 {resp.status}: {error_text}")
                
                except Exception as e:
                    last_error = str(e)
                    self.stats['retry_count'] += 1
                    
                    if attempt < self.max_retries - 1:
                        wait_time = self.request_interval * (attempt + 1)
                        await asyncio.sleep(wait_time)
            
            # 所有重试失败
            elapsed = time.time() - start_time
            self.stats['failed'] += 1
            self._update_progress(index, success=False)
            
            return ImageGenerationResult(
                index=index,
                path=None,
                success=False,
                time=elapsed,
                error=last_error
            )
    
    def _update_progress(self, index: int, success: bool = True, 
                        time: float = 0, cached: bool = False):
        """更新进度显示"""
        self.completed += 1
        
        if cached:
            status = "💾 缓存"
        elif success:
            status = f"✅ 成功 ({time:.1f}s)"
        else:
            status = "❌ 失败"
        
        # 计算预估剩余时间
        if self.completed > 0 and self.stats['total_time'] > 0:
            avg_time = self.stats['total_time'] / max(1, self.stats['success'])
            remaining = self.stats['total'] - self.completed
            eta = avg_time * remaining / max(1, self.max_concurrent)
            eta_str = f" | 预计剩余: {eta/60:.1f}分钟"
        else:
            eta_str = ""
        
        print(f"[{self.completed}/{self.stats['total']}] 图片 {index}: {status}{eta_str}")
    
    async def generate_batch(self,
                            prompts: List[str],
                            save_path: Path) -> List[Optional[str]]:
        """
        批量并行生成图片
        
        Returns:
            图片路径列表（失败的为None）
        """
        self.stats['total'] = len(prompts)
        self.completed = 0
        
        print(f"\n{'='*70}")
        print(f"🚀 并行图像生成（{self.max_concurrent}个并发）")
        print(f"{'='*70}")
        print(f"总数: {len(prompts)}")
        print(f"缓存: {'启用' if self.cache_dir else '禁用'}")
        print(f"{'='*70}\n")
        
        save_path.mkdir(parents=True, exist_ok=True)
        
        # 创建信号量控制并发
        semaphore = asyncio.Semaphore(self.max_concurrent)
        
        async with aiohttp.ClientSession() as session:
            # 创建所有任务
            tasks = [
                self._generate_single(session, semaphore, prompt, save_path, i + 1)
                for i, prompt in enumerate(prompts)
            ]
            
            # 并行执行
            results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 处理结果
        image_paths = []
        for idx, result in enumerate(results, 1):
            if isinstance(result, Exception):
                print(f"⚠️  图片 {idx} 任务异常: {result}")
                # 🔥 修复：为失败的图片创建占位符路径
                error_path = save_path / f"p{idx:04d}_error.png"
                image_paths.append(str(error_path))
            elif result.path is None:
                # 🔥 修复：生成失败时创建占位符路径
                error_path = save_path / f"p{result.index:04d}_error.png"
                image_paths.append(str(error_path))
            else:
                image_paths.append(result.path)
        
        # 打印统计
        self._print_stats()
        
        return image_paths
    
    def _print_stats(self):
        """打印统计信息"""
        print(f"\n{'='*70}")
        print("📊 生成统计")
        print(f"{'='*70}")
        print(f"总数: {self.stats['total']}")
        print(f"成功: {self.stats['success']} ({self.stats['success']/max(1,self.stats['total'])*100:.1f}%)")
        print(f"失败: {self.stats['failed']}")
        print(f"缓存命中: {self.stats['cached']}")
        print(f"重试次数: {self.stats['retry_count']}")
        
        if self.stats['success'] > 0:
            avg = self.stats['total_time'] / self.stats['success']
            total_min = self.stats['total_time'] / 60
            print(f"平均用时: {avg:.1f}秒/张")
            print(f"总耗时: {total_min:.1f}分钟")
            
            # 对比顺序生成
            sequential_time = avg * self.stats['total']
            speedup = sequential_time / max(1, self.stats['total_time'])
            print(f"提速: {speedup:.1f}x")
        
        print(f"{'='*70}\n")


# ==================== 集成方法 ====================

def optimized_stage_5_generate_images(self, prompts: list, pages: list) -> list:
    """
    替换 run_interactive.py 中的 stage_5_generate_images 方法
    
    在 InteractiveVideoGenerator 类中使用此方法
    """
    self.interactive.print_header("🖼️ 阶段5: 生成图像（优化版）")
    
    self.interactive.print_info(f"准备并行生成 {len(prompts)} 张图片...")
    
    if not self.interactive.get_yes_no("开始生成图像？（使用并行加速）", True):
        self.interactive.print_warning("跳过图像生成")
        return []
    
    # 获取配置
    image_config = self.config.get('image_generation', {})
    params = image_config.get('params', {})
    
    # 获取API密钥
    api_key = params.get('api_key')
    if not api_key:
        self.interactive.print_error("未配置 api_key")
        return []
    
    # 准备路径
    image_dir = self.story_dir / "image"
    cache_dir = self.story_dir / ".cache"
    
    # 优化提示词
    from mm_story_agent.base import init_tool_instance
    image_agent = init_tool_instance(image_config)
    style_name = params.get('style_name', 'Storybook')
    
    optimized_prompts = []
    for prompt in prompts:
        opt_prompt = image_agent.optimize_prompt(prompt, style_name)
        optimized_prompts.append(opt_prompt)
    
    # 创建优化生成器
    generator = OptimizedImageGenerator(
        api_key=api_key,
        max_concurrent=params.get('max_concurrent', 4),
        request_interval=params.get('request_interval', 1.0),
        cache_dir=cache_dir if params.get('enable_cache', True) else None,
        max_retries=params.get('max_retries', 3)
    )
    
    # 并行生成
    images = asyncio.run(
        generator.generate_batch(optimized_prompts, image_dir)
    )

    if not images or all(img is None for img in images):
        self.interactive.print_warning("返回列表为空，尝试从磁盘读取...")
        images = []
        for i in range(1, len(prompts) + 1):
            expected_path = image_dir / f"p{i:04d}.png"
            if expected_path.exists():
                images.append(str(expected_path))
                self.interactive.print_success(f"✅ 找到图片: {expected_path.name}")
            else:
                images.append(None)
                self.interactive.print_warning(f"⚠️  未找到图片: p{i:04d}.png")
    
    # 统计成功率
    success_count = sum(1 for img in images if img is not None)
    self.interactive.print_success(
        f"图像生成完成！成功 {success_count}/{len(images)} 张"
    )
    
    # 🔥 调试输出
    print(f"\n调试信息:")
    print(f"  image_paths 类型: {type(images)}")
    print(f"  image_paths 长度: {len(images)}")
    print(f"  前3个元素: {images[:3]}")
    
    # 保存进度
    self.save_progress({
        'pages': pages,
        'prompts': prompts,
        'images': images
    })
    
    # 统计成功率
    success_count = sum(1 for img in images if img is not None)
    self.interactive.print_success(
        f"图像生成完成！成功 {success_count}/{len(images)} 张"
    )
    
    return images


# ==================== 使用示例 ====================

if __name__ == "__main__":
    """
    独立测试脚本
    """
    import sys
    from pathlib import Path
    
    # 测试参数
    test_prompts = [
        "A magical forest with glowing mushrooms, fantasy style, high quality",
        "A brave knight in shining armor, medieval style, detailed",
        "A majestic dragon flying over mountains, epic style, dramatic lighting"
    ] * 3  # 9张图测试
    
    test_dir = Path("test_output")
    cache_dir = test_dir / ".cache"
    
    # 从环境变量或命令行获取API key
    api_key = sys.argv[1] if len(sys.argv) > 1 else input("请输入 DashScope API Key: ")
    
    # 创建生成器
    generator = OptimizedImageGenerator(
        api_key=api_key,
        max_concurrent=4,
        cache_dir=cache_dir
    )
    
    # 生成
    print("开始测试并行生成...")
    images = asyncio.run(generator.generate_batch(test_prompts, test_dir))
    
    print("\n生成结果:")
    for i, img in enumerate(images, 1):
        print(f"{i}. {img if img else '失败'}")